import numpy as np
import cv2
import viser
import time
from pathlib import Path
from typing import Union, List, Tuple, Optional

class PointmapVisualizer:
    def __init__(self, server_port: int = 8080):
        """
        初始化Viser点云可视化器
        
        Args:
            server_port: Viser服务器端口
        """
        self.server = viser.ViserServer(port=server_port)
        self.pointmaps = []
        self.video_frames = []
        
    def load_video(self, video_path: Union[str, Path]) -> List[np.ndarray]:
        """
        加载视频文件
        
        Args:
            video_path: 视频文件路径
            
        Returns:
            视频帧列表
        """
        cap = cv2.VideoCapture(str(video_path))
        frames = []
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            # 转换BGR到RGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
            
        cap.release()
        print(f"加载了 {len(frames)} 帧视频")
        return frames
    
    def load_pointmaps(self, pointmap_path: Union[str, Path, List]) -> List[np.ndarray]:
        """
        加载pointmap序列
        
        Args:
            pointmap_path: pointmap文件路径或路径列表，或numpy数组列表
            
        Returns:
            pointmap列表，每个元素shape为(H, W, 3)
        """
        data = np.load(pointmap_path, allow_pickle=True)
        point_map, mask = data['point_map'], data['mask']
        return point_map
    
    def pointmap_to_pointcloud(self, pointmap: np.ndarray, frame: np.ndarray, 
                              valid_depth_threshold: float = 0.1,
                              downsample_factor: int = 1) -> Tuple[np.ndarray, np.ndarray]:
        """
        将pointmap和对应的视频帧转换为点云
        
        Args:
            pointmap: 形状为(H, W, 3)的pointmap
            frame: 形状为(H, W, 3)的视频帧
            valid_depth_threshold: 有效深度阈值，过滤掉深度过小的点
            downsample_factor: 下采样因子，1表示不下采样，2表示每2个像素取1个
            
        Returns:
            points: (N, 3) 3D点坐标
            colors: (N, 3) 对应的RGB颜色 (0-1范围)
        """
        H, W = pointmap.shape[:2]
        
        # 下采样处理
        if downsample_factor > 1:
            # 使用步长进行下采样
            pointmap = pointmap[::downsample_factor, ::downsample_factor]
            frame = frame[::downsample_factor, ::downsample_factor]
        
        # 展平pointmap和frame
        points_3d = pointmap.reshape(-1, 3)  # (H*W, 3)
        colors_rgb = frame.reshape(-1, 3)    # (H*W, 3)
        
        # 过滤有效点（深度大于阈值的点）
        valid_mask = np.linalg.norm(points_3d, axis=1) > valid_depth_threshold
        
        points = points_3d[valid_mask]
        colors = colors_rgb[valid_mask] / 255.0  # 归一化到0-1
        
        return points, colors
    
    def visualize_sequence(self, video_path: Union[str, Path], 
                          pointmap_data: Union[str, Path, List],
                          frame_interval: float = 0.1,
                          point_size: float = 0.02):
        """
        可视化pointmap序列
        
        Args:
            video_path: 视频文件路径
            pointmap_data: pointmap数据（文件路径或numpy数组列表）
            frame_interval: 帧间隔时间（秒）
            point_size: 点的大小
        """
        # 加载数据
        self.video_frames = self.load_video(video_path)[::5] # [t, h, w, 3]
        self.pointmaps = self.load_pointmaps(pointmap_data)[::5] # [t, h, w, 3]
        print(f"读入视频与pointmap: {self.video_frames[0].shape, self.pointmaps.shape}")
        
        # 确保帧数匹配
        min_frames = min(len(self.video_frames), len(self.pointmaps))
        print(f"将可视化 {min_frames} 帧")
        
        # 添加GUI控件
        with self.server.gui.add_folder("播放控制"):
            frame_slider = self.server.gui.add_slider(
                "Frame", min=0, max=min_frames-1, step=1, initial_value=0
            )
            play_button = self.server.gui.add_button("Play/Pause")
            
            # FPS控制
            fps_slider = self.server.gui.add_slider(
                "FPS", min=1, max=60, step=1, initial_value=int(1/frame_interval)
            )
            
        with self.server.gui.add_folder("渲染设置"):
            point_size_slider = self.server.gui.add_slider(
                "Point Size", min=0.001, max=0.1, step=0.001, initial_value=point_size
            )
            
            # 下采样控制
            downsample_slider = self.server.gui.add_slider(
                "Downsample Factor", min=1, max=8, step=1, initial_value=1
            )
            
            # 深度阈值控制
            depth_threshold_slider = self.server.gui.add_slider(
                "Depth Threshold", min=0.01, max=1.0, step=0.01, initial_value=0.1
            )
            
        # with self.server.gui.add_folder("相机控制"):
            # # 相机预设位置按钮
            # camera_front_button = self.server.gui.add_button("前视图")
            # camera_top_button = self.server.gui.add_button("俯视图")
            # camera_side_button = self.server.gui.add_button("侧视图")
            # camera_reset_button = self.server.gui.add_button("重置相机")
            
            # # 相机参数调节
            # camera_distance_slider = self.server.gui.add_slider(
            #     "Camera Distance", min=0.5, max=10.0, step=0.1, initial_value=3.0
            # )
            # camera_height_slider = self.server.gui.add_slider(
            #     "Camera Height", min=-2.0, max=2.0, step=0.1, initial_value=0.0
            # )
            
        # 播放状态
        playing = False
        current_frame = 0
        last_time = time.time()
        
        def update_pointcloud(frame_idx: int):
            """更新点云显示"""
            if frame_idx >= min_frames:
                return
                
            # 转换为点云，使用当前的下采样和深度阈值设置
            points, colors = self.pointmap_to_pointcloud(
                self.pointmaps[frame_idx], 
                self.video_frames[frame_idx],
                valid_depth_threshold=depth_threshold_slider.value,
                downsample_factor=downsample_slider.value
            )
            
            # 更新viser中的点云
            self.server.scene.add_point_cloud(
                name="pointcloud",
                points=points,
                colors=colors,
                point_size=point_size_slider.value,
            )
            
            print(f"Frame {frame_idx}: {len(points)} points (downsample: {downsample_slider.value}x)")
            
        def set_camera_view(position: Tuple[float, float, float], 
                           target: Tuple[float, float, float] = (0, 0, 0)):
            """设置相机视角"""
            try:
                # 使用 viser 的相机控制方法
                # 计算 look_at 向量
                look_at = np.array(target) - np.array(position)
                look_at = look_at / np.linalg.norm(look_at)
                
                # 计算 up 向量
                up = np.array([0, 0, 1])
                
                # 使用 viser 的相机设置方法
                # 注意：不同版本的 viser 可能有不同的 API
                if hasattr(self.server, 'camera'):
                    self.server.camera.position = position
                else:
                    # 如果没有直接的相机控制，我们可以通过添加一个隐形的相机标记来提示用户
                    self.server.scene.add_frame(
                        name="camera_hint",
                        axes_length=0.3,
                        axes_radius=0.01,
                    )
                    # 移动到指定位置（这会影响整个场景的视角）
                    print(f"建议相机位置: {position}, 目标: {target}")
                    
            except Exception as e:
                print(f"相机设置失败: {e}")
                print("请手动调整视角")
        
        # 回调函数
        @frame_slider.on_update
        def _(_):
            nonlocal current_frame
            current_frame = frame_slider.value
            update_pointcloud(current_frame)
        
        @play_button.on_click
        def _(_):
            nonlocal playing
            playing = not playing
        
        @point_size_slider.on_update
        def _(_):
            # 重新渲染当前帧以更新点大小
            update_pointcloud(current_frame)
            
        @downsample_slider.on_update
        def _(_):
            # 重新渲染当前帧以更新下采样
            update_pointcloud(current_frame)
            
        @depth_threshold_slider.on_update
        def _(_):
            # 重新渲染当前帧以更新深度阈值
            update_pointcloud(current_frame)
            
        @fps_slider.on_update
        def _(_):
            # FPS变化时不需要特别处理，播放循环会自动使用新的FPS值
            pass
            
        # # 相机控制回调
        # @camera_front_button.on_click
        # def _(_):
        #     distance = camera_distance_slider.value
        #     height = camera_height_slider.value
        #     set_camera_view((0, -distance, height), (0, 0, 0))
        #     print(f"切换到前视图 - 建议位置: (0, {-distance}, {height})")
            
        # @camera_top_button.on_click
        # def _(_):
        #     distance = camera_distance_slider.value
        #     set_camera_view((0, 0, distance), (0, 0, 0))
        #     print(f"切换到俯视图 - 建议位置: (0, 0, {distance})")
            
        # @camera_side_button.on_click
        # def _(_):
        #     distance = camera_distance_slider.value
        #     height = camera_height_slider.value
        #     set_camera_view((distance, 0, height), (0, 0, 0))
        #     print(f"切换到侧视图 - 建议位置: ({distance}, 0, {height})")
            
        # @camera_reset_button.on_click
        # def _(_):
        #     distance = camera_distance_slider.value
        #     height = camera_height_slider.value
        #     set_camera_view((distance * 0.7, -distance * 0.7, height + 1), (0, 0, 0))
        #     print(f"重置视角 - 建议位置: ({distance * 0.7:.1f}, {-distance * 0.7:.1f}, {height + 1})")
        #     # 清除所有可能的指示器（如果有的话）
        #     try:
        #         self.server.scene.remove("view_indicator_front")
        #         self.server.scene.remove("view_indicator_top") 
        #         self.server.scene.remove("view_indicator_side")
        #         self.server.scene.remove("camera_hint")
        #     except:
        #         pass
                
        # @camera_distance_slider.on_update
        # def _(_):
        #     # 显示距离变化的提示
        #     print(f"相机距离调整为: {camera_distance_slider.value}")
                
        # @camera_height_slider.on_update
        # def _(_):
        #     # 显示高度变化的提示
        #     print(f"相机高度调整为: {camera_height_slider.value}")
        
        # 初始显示第一帧
        update_pointcloud(0)
        
        # 添加坐标轴作为参考 - 缩小尺寸
        self.server.scene.add_frame(
            name="world_axis",
            axes_length=0.3,
            axes_radius=0.01,
            position=(0, 0, 0)
        )
        
        # 播放循环
        try:
            while True:
                current_time = time.time()
                
                if playing:
                    # 根据FPS计算帧间隔
                    target_interval = 1.0 / fps_slider.value
                    
                    if current_time - last_time >= target_interval:
                        current_frame = (current_frame + 1) % min_frames
                        frame_slider.value = current_frame
                        update_pointcloud(current_frame)
                        last_time = current_time
                    
                time.sleep(0.01)  # 小的睡眠时间以避免CPU占用过高
                
        except KeyboardInterrupt:
            print("可视化已停止")

# 使用示例
def main():
    # 创建可视化器
    visualizer = PointmapVisualizer(server_port=8080)
    
    # 方式1: 从文件加载
    video_path = "examples/video1.mp4"
    pointmap_path = "outputs/video1.npz"  # pointmap文件列表
    
    # 方式2: 直接使用numpy数组
    # pointmaps = [pointmap_array_0, pointmap_array_1, ...]  # 已经加载的pointmap数组列表
    
    try:
        # 开始可视化
        print("启动Viser服务器...")
        print("请在浏览器中打开 http://localhost:8080 查看可视化")
        
        visualizer.visualize_sequence(
            video_path=video_path,
            pointmap_data=pointmap_path,  # 或者直接传入pointmaps数组列表
            frame_interval=0.1,  # 100ms间隔（初始值，可通过FPS滑块调整）
            point_size=0.02
        )
        
    except Exception as e:
        print(f"错误: {e}")


if __name__ == "__main__":
    main()